import os
import json
from datetime import datetime
from typing import Dict, Any

from rainforeLearn.gomoku.v2.train.constants.train_stats_printer_constants import TrainingStatsPrinterConstants
from rainforeLearn.gomoku.v2.train.statistics.tracker.episode_game_result_tracker import EpisodeGameResultTracker
from rainforeLearn.gomoku.v2.train.statistics.tracker.reward_penalty_tracker import RewardPenaltyTracker
from rainforeLearn.gomoku.v2.train.statistics.tracker.threat_defense_tracker import ThreatDefenseTracker


class LogGenerator:
    """训练日志生成器 - 负责生成和保存训练日志"""
    
    def __init__(self, game_tracker: EpisodeGameResultTracker, 
                 reward_penalty_tracker: RewardPenaltyTracker,
                 threat_stats: ThreatDefenseTracker):
        self.game_tracker = game_tracker
        self.reward_penalty_tracker = reward_penalty_tracker
        self.threat_stats = threat_stats
    
    def save_training_log(self, episode: int, checkpoint_dir: str = 'checkpoints') -> None:
        """保存训练日志"""
        os.makedirs(checkpoint_dir, exist_ok=True)

        log_data = self._create_log_data(episode)
        log_path = os.path.join(checkpoint_dir,
                                f'{TrainingStatsPrinterConstants.LOG_FILE_PREFIX}{episode}.json')

        with open(log_path, 'w', encoding='utf-8') as f:
            json.dump(log_data, f, indent=2, ensure_ascii=False)

        print(f"📝 训练日志已保存: {log_path}")

    def _create_log_data(self, episode: int) -> Dict[str, Any]:
        """创建日志数据"""
        total_games = self.game_tracker.get_total_games()

        return {
            'episode': episode,
            'timestamp': datetime.now().isoformat(),
            'game_stats': {
                'total_games': total_games,
                'wins': self.game_tracker.total_wins,
                'losses': self.game_tracker.total_losses,
                'draws': self.game_tracker.total_draws,
                'win_rate': self.game_tracker.get_total_win_rate()
            },
            'reward_penalty_stats': {
                'total_batch_bonus': self.reward_penalty_tracker.total_batch_bonus,
                'total_batch_penalties': self.reward_penalty_tracker.total_batch_penalties,
                'net_bonus': self.reward_penalty_tracker.get_net_bonus(),
                'win_rate_bonus_history': self.reward_penalty_tracker.win_rate_bonus_history[
                                          -TrainingStatsPrinterConstants.RECENT_HISTORY_LIMIT:],
                'batch_penalty_history': self.reward_penalty_tracker.batch_penalty_history[
                                         -TrainingStatsPrinterConstants.RECENT_HISTORY_LIMIT:]
            },
            'threat_defense_stats': self.threat_stats.to_dict(),
            'recent_performance': {
                'win_rate_history': self.game_tracker.win_rate_history[-TrainingStatsPrinterConstants.RECENT_HISTORY_LIMIT:]
            }
        }